package com.kpbvip.aikf.controller;

import com.alibaba.fastjson.JSONObject;
import com.kpbvip.aikf.service.OpenRouterService;
import com.kpbvip.aikf.util.SseUtil;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

@RestController
@RequestMapping("/api/chat")
public class ChatController {

    @Autowired
    private OpenRouterService openRouterService;
    @RequestMapping(value = "/completion")
    @ResponseBody
    public String strChatCompletion(@RequestBody String prompt) {
        System.out.println("入参："+prompt);
            String val = openRouterService.getChatCompletion(prompt);
            System.out.println("出参："+val);
            return val;

    }



    @RequestMapping(value = "/completions")
    public SseEmitter streamChatCompletion(@RequestBody String prompt) {

            System.out.println("入参：" + prompt);
            JSONObject jsonObject = JSONObject.parseObject(prompt);
            String model = jsonObject.getString("model");
            if("AAA".equalsIgnoreCase(model)){
                Map<String,String> headMap= new HashMap<>();
                headMap.put("Authorization","Bearer fastgpt-zL26g2LP2tw5rxY4LGqrWjQ7LYDKh9aZWrAYKaFEbCSRZqvylTRQSvFRkqjConCkY");
                return SseUtil.send("http://192.168.170.100:3000/api/v1/chat/completions",prompt,headMap);
            }

            SseEmitter emitter = new SseEmitter(-1l);
            try {
                Flux<OpenAiApi.ChatCompletionChunk> result = openRouterService.getChatCompletionStream(prompt);
                result.subscribe(chunk -> {
                    try {
                        emitter.send(SseEmitter.event().data(chunk));
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                }, error -> emitter.completeWithError(error), emitter::complete);
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
            return emitter;

    }



}